import pygame
import sys
import argparse
import socket
import time

ZERO_FILTER = 0.07
joy_udp_sock: socket.socket = None


def test_udp(ip, udp_port):
    """
    发送 'test' 消息到指定的 UDP 端口，并检查是否收到 'okay' 响应。

    参数:
    ip (str): 目标 IP 地址。
    udp_port (int): 目标 UDP 端口。
    """
    test_message = b"test"
    expected_response = b"okay"
    timeout = 5  # 超时时间，单位为秒
    # 创建 UDP 套接字
    udp_sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    udp_sock.settimeout(timeout)
    try:
        # 记录发送时间
        start_time = time.time()
        # 发送 'test' 消息
        udp_sock.sendto(test_message, (ip, udp_port))
        print(f"Sent 'test' to UDP port {udp_port}")
        # 接收响应
        response, addr = udp_sock.recvfrom(1024)
        # 记录接收时间
        end_time = time.time()
        # 计算延迟
        delay = end_time - start_time
        if response == expected_response:
            print(f"Received 'okay' from {addr}")
            print(f"Round-trip delay: {delay:.6f} seconds")
            return True
        else:
            print(f"Unexpected response: {response}")
            return False
    except socket.timeout:
        print("No response received within the timeout period")
        return False
    finally:
        udp_sock.close()

def send_udp_message(message: str, ip, port):
    """
    发送UDP消息。

    参数:
    message (str): 要发送的消息。
    ip (str): 目标IP地址。
    port (int): 目标端口。
    """
    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    sock.sendto(message.encode(), (ip, port))
    print(f"Sent message: {message} to {ip}:{port}")

def scan_joysticks():
    pygame.init()
    pygame.joystick.init()
    joystick_count = pygame.joystick.get_count()
    joysticks = {}
    for i in range(joystick_count):
        joystick = pygame.joystick.Joystick(i)
        # joystick.init()  # Deprecated
        joysticks[i] = joystick.get_name()
    return joysticks

def udp_send_axis_speed(axis_status, ip, port):
    out_values = [0., 0., 0.]
    out_values[0] = speed_map(axis_status[1], -1, 1, 100, -100)
    out_values[1] = speed_map(axis_status[3], -1, 1, 5000, -5000)
    out_values[2] = speed_map(axis_status[4], -1, 1, 5000, -5000)
    
    joy_udp_sock.sendto(f":0 {out_values[0]:.3f}\r\n".encode(), (ip, port))
    joy_udp_sock.sendto(f":1 {out_values[1]:.3f}\r\n".encode(), (ip, port))
    joy_udp_sock.sendto(f":2 {out_values[2]:.3f}\r\n".encode(), (ip, port))

def udp_init(ip, port):
    global joy_udp_sock
    if not test_udp(ip, port):
        print(f"UDP connection {ip}:{port} failed")
        sys.exit(1)
    if joy_udp_sock is not None:
        try:
            joy_udp_sock.close()
        except Exception as e:
            print(f"Error closing UDP socket: {e}")
    joy_udp_sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    # joy_udp_sock.bind((ip, port))

def speed_map(value, from_min, from_max, to_min, to_max):
    """
    将数值从一个区间映射到另一个区间。

    参数:
    value (float): 要映射的数值。
    from_min (float): 原区间的最小值。
    from_max (float): 原区间的最大值。
    to_min (float): 目标区间的最小值。
    to_max (float): 目标区间的最大值。

    返回:
    float: 映射后的数值。
    """
    # 计算比例
    scale = (to_max - to_min) / (from_max - from_min)
    # 映射数值
    return to_min + (value - from_min) * scale

def main(joy_id, ip, port):
    pygame.init()
    joystick = pygame.joystick.Joystick(joy_id)
    clock = pygame.time.Clock()
    udp_init(ip, port)
    while True:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                return
        axis_values = [joystick.get_axis(i) for i in range(joystick.get_numaxes())]
        axis_values[-2] = (-axis_values[-2] + axis_values[-1]) / 2  # Invert the Y axis
        axis_values.pop()
        filtered_axis_values = [0 if abs(value) < ZERO_FILTER else value for value in axis_values]
        print(f"Axis values: {filtered_axis_values}")
        udp_send_axis_speed(filtered_axis_values, ip, port)
        clock.tick(10)  # 控制刷新率在10Hz

def test_joy(joy_id):
    pygame.init()
    joystick = pygame.joystick.Joystick(joy_id)
    clock = pygame.time.Clock()
    
    while True:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                return
        axis_values = [joystick.get_axis(i) for i in range(joystick.get_numaxes())]
        axis_values[-2] = (-axis_values[-2] + axis_values[-1]) / 2  # Invert the Y axis
        axis_values.pop()
        filtered_axis_values = [0 if abs(value) < ZERO_FILTER else value for value in axis_values]
        print(f"Axis values: {filtered_axis_values}")
        clock.tick(10)  # 控制刷新率在10Hz

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="UDP Joystick Client")
    parser.add_argument('--scan', action='store_true', help='Scan for available joysticks')
    parser.add_argument('--joy_id', type=int, default=0, help='Joystick ID to use')
    parser.add_argument('--ip', type=str, default="0.0.0.0", help='IP address to send UDP messages to')
    parser.add_argument('--port', type=int, default=8080, help='UDP port to send messages to')
    parser.add_argument('--testjoy', action='store_true', help='Test joystick without sending UDP messages')
    args = parser.parse_args()

    if args.scan:
        joysticks = scan_joysticks()
        print(f"Number of joysticks: {len(joysticks)}")
        for id, name in joysticks.items():
            print(f"Joystick {id}: {name}")
    elif args.testjoy:
        test_joy(args.joy_id)
    else:
        try:
            main(args.joy_id, args.ip, args.port)
        except KeyboardInterrupt:
            print("Exiting...")
        finally:
            if joy_udp_sock:
                joy_udp_sock.close()